# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import re

system_prompt = '''# 最高级别系统提示

Sonetto，一个多 AI 智能体对话协作框架，旨在发挥各个智能体的优势，充分释放大模型在推理、规划方面的潜力。Sonetto 通过将不同的智能体组合在一起，形成一个强大的协作网络，使得每个智能体都能在其擅长的领域内发挥最大作用。
Sonetto 的设计理念是将不同的智能体组成一个团队，通过协同工作来解决复杂的问题。每个智能体都有其独特的能力和知识领域，Sonetto 通过智能体之间的对话和协作，实现信息共享和知识融合。

你是用户基于 Sonetto 框架搭建的团队中的一员。你可以与其他成员进行对话，获取信息和知识，并通过协作来解决问题。

## Sonetto 通信协议
为确保高效协作，除函数调用的通信情形外，你必须严格遵循以下 XML 通信协议和规范：

### 1. 消息格式
除了函数调用情形外，你的回复必须使用以下XML结构，保证消息能被XLM解析器正确解析（无需严格缩进）。元素名称必须和下面一致。

<sonetto-message>
  <header>
    <to>发送者需要提供该标签，消息接收方 id 列表，每个 id 以空格分隔，例：chunker spliter checker 所有接收方都会收到该消息，但只有第一个接收方会提供响应</to>
    <from>接收者收到的消息会包含该标签，消息发送方 id</from>
    <type>协作/通知/信息/请求/响应/命令/结果/(其它可以简单描述信息作用的类型均可)</type>
  </header>
  <content>
    这里是具体的消息内容
  </content>
</sonetto-message>

### 2. 消息路由机制
路由包含两个选择：

- 单播：消息发送方指定接收方，接收方会收到该消息，其他智能体不会收到该消息。
- 组播：消息发送方指定接收方组，所有接收方都会收到该消息，但只有第一个接收方会提供响应。

### 3. 最佳实践

1. 必须提供接收方 id 列表，接收方 id 列表可以是单个 id，也可以是多个 id 以空格分隔。确保接收方 id 填写正确。
2. 必须提供消息类型，消息类型可以是协作/信息/请求/响应/命令/结果/(其它可以简单描述信息作用的类型均可)。
3. 必须提供消息内容，消息内容承载了真正需要传递的信息。
4. 明确你的团队职责。
5. 正确选择消息接收方，避免打扰无关团队成员，降低团队协作效率。
6. 函数调用情形无需遵循以上 XML 通信协议和规范。
7. 所有对话内容都写在 content 标签中. sonetto-message 标签以外的内容会被忽略。
8. 一次回答必须有且仅能包含一个 sonetto-message 标签，不能嵌套。

===

## 团队成员
id: user, 职责: 布置总任务。团队的终极目标即通过协作完成用户布置的任务。
{members}

## 你在团队中的角色
{role}。

明确你的职责，积极参与团队协作，完成用户布置的任务。

===

下面是一些额外的指导：

'''


def wrap_content(content: str,
                 _from: str = '',
                 to: str = '',
                 message_type: str = '') -> str:
    return f'''<sonetto-message>
  <header>
    <to>{to}</to>
    <from>{_from}</from>
    <type>{message_type}</type>
    <comment>请严格遵循 Sonetto 通信协议格式，严格遵循系统提示的指导</comment>
  </header>
  <content>
    {content}
  </content>
</sonetto-message>'''


def unwrap_message(message: str) -> dict | None:
    """
    解析sonetto-message格式的XML字符串，提取header和content内容
    
    Args:
        message: 包含sonetto-message的XML字符串
        
    Returns:
        dict: 包含解析后的消息内容的字典
    """

    # 清理模型深度思考输出

    # 清理可能存在的多余空白和格式问题
    message = message.strip()

    # 检查是否为有效的sonetto-message
    # if not message.startswith('<sonetto-message>') or not message.endswith('</sonetto-message>'):
    #     return None
    # else:
    #     print('XML parse correctly')

    # try:
    #     # 解析XML
    #     root = ET.fromstring(message)
    #     sonetto_message = root.find('sonetto-message')
    #     if sonetto_message is None:
    #         return None
    #     else:
    #         print('XML parse correctly')

    #     # 提取header信息
    #     header = root.find('header')
    #     if header is None:
    #         return None

    #     to = header.find('to').text.strip() if header.find(
    #         'to') is not None else ""
    #     from_agent = header.find('from').text.strip() if header.find(
    #         'from') is not None else ""
    #     message_type = header.find('type').text.strip() if header.find(
    #         'type') is not None else ""

    #     # 提取content内容
    #     content = root.find('content').text.strip() if root.find(
    #         'content') is not None else ""

    #     # 返回解析结果
    #     return {
    #         'to': to,
    #         'from': from_agent,
    #         'message_type': message_type,
    #         'content': content
    #     }

    # except Exception as e:
        # 如果XML解析失败，尝试使用正则表达式
    try:
        sonetto_message = re.search(r'<sonetto-message>(.*?)</sonetto-message>', message, re.DOTALL)
        if not sonetto_message:
            return None
        to_pattern = re.compile(r'<to>(.*?)</to>', re.DOTALL)
        from_pattern = re.compile(r'<from>(.*?)</from>', re.DOTALL)
        type_pattern = re.compile(r'<type>(.*?)</type>', re.DOTALL)
        content_pattern = re.compile(r'<content>\s*(.*?)\s*</content>',
                                        re.DOTALL)

        to = to_pattern.search(message).group(
            1).strip() if to_pattern.search(message) else ""
        if not to:
            return None
        from_agent = from_pattern.search(message).group(
            1).strip() if from_pattern.search(message) else ""
        message_type = type_pattern.search(message).group(
            1).strip() if type_pattern.search(message) else ""
        content = content_pattern.search(message).group(
            1).strip() if content_pattern.search(message) else ""

        return {
            'to': to,
            'from': from_agent,
            'message_type': message_type,
            'content': content
        }

    except Exception as e2:
        # return {'error': f'Failed to parse message: {str(e2)}'}
        return None
